//
// Ephi - simulation of magnetic fields and particles
// Copyright (C) 2007 Indrek Mandre <indrek(at)mare.ee>
// For more information please see http://www.mare.ee/indrek/ephi/
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>

#include "ephi.hpp"
#include "potential.hpp"

#define LOOPS 10

class PolywellOctreeBuilder : public OctreeBuilder
{
public:
  PolywellOctreeBuilder (PMap& pmap, const vect3d& c1, const vect3d& c2)
    : OctreeBuilder (c1, c2, 1, 2, 100, 0.01), pmap(pmap)
  {
    addLimit (vect3d(0, 1, 0));
    addLimit (vect3d(1, 0, -1).normal());
    addLimit (vect3d(0, -1, 1).normal());
  }

  void addValues (const vect3d& pos, CubicScalar *scalars, CubicVector *vectors)
  {
    pmap.addCubicValues (pos, scalars[0], vectors[0], vectors[1]);
  }

  bool inContact (const vect3d& pos)
  {
    return false;
  }

private:
  PMap &pmap;
};

static void usage (const char *name)
{
  fprintf (stderr, "%s: %s [-m pointsystem|linesystem|chargeflow] [-q] [-x] [-s] [-g] [-8] [-f] [-r resolution] [-y]\n", name, name);
  exit(-1);
}

int main (int argc, char *argv[])
{
  prec_t resolution = 0.018;
  bool quick = false;
  bool cross = false;
  bool gun = false;
  bool snapshots = false;
  bool at8 = false;
  bool fcage = false;
  bool symm = false;
  PField::emethod_t method = PField::POINTSYSTEM;

  for ( int i = 1; i < argc; i++ )
    {
      if ( argv[i][0] != '-' )
          usage(argv[0]);
      switch (argv[i][1])
        {
          case 'm': {
            if ( ++i == argc )
                usage(argv[0]);
            if ( !strcmp (argv[i], "pointsystem") )
                method = PField::POINTSYSTEM;
            else if ( !strcmp (argv[i], "linesystem") )
                method = PField::LINESYSTEM;
            else if ( !strcmp (argv[i], "chargeflow") )
                method = PField::CHARGEFLOW;
            else
                usage(argv[0]);
            break;
          }
          case 'f': {
            fcage = true;
            break;
          }
          case '8': {
            at8 = true;
            break;
          }
          case 'g': {
            gun = true;
            break;
          }
          case 's': {
            snapshots = true;
            break;
          }
          case 'x': {
            cross = true;
            break;
          }
          case 'r': {
            if ( ++i == argc )
                usage(argv[0]);
            resolution = atof(argv[i]);
            break;
          }
          case 'q': {
            quick = true;
            break;
          }
          case 'y': {
            symm = true;
            break;
          }
          default: {
            usage(argv[0]);
            break;
          }
        }
    }

  prec_t cali_pot_min = 0;
  prec_t cali_pot_max = 0;
  prec_t cali_pot_add = 0;
  prec_t cali_e_min = 0;
  prec_t cali_e_max = 0;

  prec_t potential = 10000;
  prec_t r = 0.15;
  prec_t wr = 0.035;
  prec_t spacing = 0.01;
  prec_t rr = r + (2 * wr + spacing) / PREC_SQRT2;
  prec_t qinc = -1.5e-8;
  prec_t egun;
  vect3d egunpos;
  if ( at8 )
    {
      egunpos = vect3d(-1, 1, 1).normal() * (rr + r);
      egun = egunpos.y;
    }
  else
    {
      egunpos = vect3d(-(rr + r), 0, 0);
      egun = -egunpos.x;
    }

  FILE *csvfp = fopen ("polyp.csv", "w");
  if ( !csvfp )
    {
      fprintf (stderr, "%s: unable to open %s for writing\n", argv[0], "polyp.csv");
      exit(-1);
    }
  if ( gun )
      fprintf (csvfp, "Coil voltage[V],Space charge[C],Center potential[V],Potential well[V],Electron acceleration[V]\n");
  else
      fprintf (csvfp, "Coil voltage[V],Space charge[C],Coil charge[C],Center potential[V],Potential well[V],Electron acceleration[V]\n");

  for ( size_t i = LOOPS; i > 0; i-- )
    {
      char buf[32];
      sprintf (buf, "polyp%02u", (unsigned int)i);
      std::string prefix = buf;
      PField pf(method, resolution);
      prec_t q = (i - 1) * qinc;
#if 1
#if 1
      pf.addTorus (vect3d(-rr, 0, 0), vect3d(1, 0, 0), r, wr, potential);
      //pf.set_no_cdist (true);
      pf.addTorus (vect3d(rr, 0, 0), vect3d(-1, 0, 0), r, wr, potential);
      //pf.set_no_cdist (false);
      pf.addTorus (vect3d(0, -rr, 0), vect3d(0, 1, 0), r, wr, potential);
      pf.addTorus (vect3d(0, rr, 0), vect3d(0, -1, 0), r, wr, potential);
      pf.addTorus (vect3d(0, 0, rr), vect3d(0, 0, -1), r, wr, potential);
      pf.addTorus (vect3d(0, 0, -rr), vect3d(0, 0, 1), r, wr, potential);
#endif

      std::vector<vect3d> splanes;
      splanes.push_back (vect3d(0, 1, 0));
      splanes.push_back (vect3d(1, 0, -1).normal());
      splanes.push_back (vect3d(0, -1, 1).normal());

      if ( gun )
        {
          prec_t c = r / 4;
          if ( c < resolution )
              c = resolution;
          prec_t offset = egun + c + c / 4;

          prec_t ores = pf.setResolution (resolution * 2);
          if ( at8 )
            {
              if ( symm )
                {
                  pf.setResolution (ores);
                  vect3d n(1, 1, 1);
                  n.normalize();
                  vect3d mid(offset, offset, offset);
                  vect3d pos1 = mid + n * c;
                  vect3d pos2 = mid - n * c;
                  vect3d s1 = (splanes[1] % n).normal();
                  vect3d pos3 = mid + s1 * c;
                  vect3d s2 = (splanes[2] % n).normal();
                  vect3d pos4 = mid - s2 * c;

                  std::vector<PPoint*> spoints;
                  std::vector<PLink*> slinks;

                  prec_t pratio = 0.25;
                  spoints.push_back (&pf.addPoint (pos1, pos1 + pratio * (mid - pos1), 0));
                  spoints.push_back (&pf.addPoint (pos2, pos2 + pratio * (mid - pos2), 0));
                  spoints.push_back (&pf.addPoint (pos3, pos3 + pratio * (mid - pos3), 0));
                  spoints.push_back (&pf.addPoint (pos4, pos4 + pratio * (mid - pos4), 0));

                  slinks.push_back (&pf.connect (*spoints[0], *spoints[2]));
                  slinks.push_back (&pf.connect (*spoints[0], *spoints[3]));
                  slinks.push_back (&pf.connect (*spoints[1], *spoints[2]));
                  slinks.push_back (&pf.connect (*spoints[1], *spoints[3]));
                  slinks.push_back (&pf.connect (*spoints[2], *spoints[3]));

                  pf.splitSphericalTriangles (mid, c, spoints, slinks, 0, false, pratio);
                }
              else
                {
                  pf.addSphere (vect3d(-offset, offset, offset), c, 0);
                  pf.addSphere (vect3d(-offset, offset, -offset), c, 0);
                  pf.addSphere (vect3d(-offset, -offset, offset), c, 0);
                  pf.addSphere (vect3d(-offset, -offset, -offset), c, 0);
                  pf.addSphere (vect3d(offset, -offset, offset), c, 0);
                  pf.addSphere (vect3d(offset, -offset, -offset), c, 0);
                  pf.addSphere (vect3d(offset, offset, offset), c, 0);
                  pf.addSphere (vect3d(offset, offset, -offset), c, 0);
                }
            }
          else
            {
              pf.addSphere (vect3d(-offset, 0, 0), c, 0);
              pf.addSphere (vect3d(offset, 0, 0), c, 0);
              pf.addSphere (vect3d(0, -offset, 0), c, 0);
              pf.addSphere (vect3d(0, offset, 0), c, 0);
              pf.addSphere (vect3d(0, 0, -offset), c, 0);
              pf.addSphere (vect3d(0, 0, offset), c, 0);
            }
          pf.setResolution (ores);
        }

      prec_t cage_radius = egunpos.length() + r + r/2;

      if ( fcage )
        {
          pf.set_no_cdist (true);
          prec_t ores = pf.setResolution (cage_radius / 10.0);
          if ( symm )
            {
              vect3d n1 = (splanes[1] % splanes[0]).normal();
              vect3d n2 = (splanes[0] % splanes[2]).normal();
              vect3d n3 = (splanes[2] % splanes[1]).normal();
              pf.addSphereSegment (cage_radius, n1, n2, n3, 0, true, 2);
            }
          else
            {
              pf.addSphere (vect3d(0, 0, 0), cage_radius, 0, true);
            }
          pf.setResolution(ores);
          pf.set_no_cdist (false);
        }

      pf.setUniformCharge (vect3d(0, 0, 0), q, 0.1);

      if ( symm && (method == PField::LINESYSTEM || method == PField::POINTSYSTEM) )
        {
          std::vector<transf3d> transforms;
          transforms.reserve (48);
          const transf3d_pair *tr = get_cube_polywell_transforms48 ();
          for ( size_t i = 0; i < 48; i++ )
              transforms.push_back (tr[i].f);
          pf.setTransforms (splanes, transforms);
        }

      if ( !pf.equalize() )
        {
          fprintf (stderr, "equalizing failed :(\n");
          abort();
        }
      pf.save (prefix + "_model.pot");
#else
      pf.load (prefix + "_model.pot");
#endif

      prec_t gunaccel = 0;
      vect3d spos = egunpos;
      vect3d gdir = -egunpos / 1000.0;
      for ( size_t k = 0; k < 1000; k++ )
        {
          prec_t pot = pf.getPotential (spos);
          if ( pot < gunaccel )
              break;
          gunaccel = pot;
          spos += gdir;
        }

      vect3d pot_pos(0, 0, 0);
      prec_t pot_well = pf.getPotential (pot_pos);
      while (1)
        {
          pot_pos = pot_pos + (r / 1000.0) * vect3d(1, 0, 0);
          prec_t pot = pf.getPotential (pot_pos);
          if ( pot < pot_well ) break;
          pot_well = pot;
        }

      std::string str = pf.getChargeConcentrationScript ();
      FILE *fp = fopen ((prefix + "_charge.pov").c_str(), "w");
      fprintf (fp, "camera { location <3, 3, 2>*0.3 look_at <0, 0, 0> sky <0, 0, 1> right x*image_width/image_height }\n");
      fprintf (fp, "#include \"scene.inc\"\n");
      fprintf (fp, "%s", str.c_str());
      fclose (fp);

      str = pf.getPovrayScript ();
      fp = fopen ((prefix + "_polywell.pov").c_str(), "w");
      fprintf (fp, "camera { location <3, 3, 2>*0.3 look_at <0, 0, 0> sky <0, 0, 1> right x*image_width/image_height }\n");
      fprintf (fp, "#include \"scene.inc\"\n");
      fprintf (fp, "%s", str.c_str());
      fclose (fp);

      fp = fopen ((prefix + "_info.txt").c_str(), "w");
      prec_t added_charge = pf.getAddedCharge();
      fprintf (fp, "Coil voltage = %.2f V\n", prec2double(potential));
      if ( !gun )
        {
          fprintf (fp, "Coil charge = %.5e C\n", prec2double(added_charge));
          fprintf (fp, "Coil Capacitance = %.3f pF\n", prec2double(added_charge / potential)*1e12);
        }
      fprintf (fp, "Space charge = %.5e C\n", prec2double(q));
      prec_t pot_mid = pf.getPotential(vect3d(0, 0, 0));
      fprintf (fp, "Potential in the middle: %.2f V\n", prec2double(pot_mid));
      fprintf (fp, "Potential well: %.2f V\n", prec2double(fabs(pot_mid - pot_well)));
      prec_t atgun;
      if ( gun )
          atgun = 0;
      else
          atgun = pf.getPotential (egunpos);
      prec_t eacc = gunaccel - atgun;
      fprintf (fp, "Electron acceleration: %.2f V\n", prec2double(eacc));
      fclose (fp);

      if ( !gun )
        {
          fprintf (csvfp, "%g,%g,%g,%g,%g,%g\n",
              prec2double (potential), prec2double(q), prec2double(added_charge),
              prec2double (pot_mid), prec2double(pot_well - pot_mid), prec2double(eacc));
        }
      else
        {
          fprintf (csvfp, "%g,%g,%g,%g,%g\n", prec2double (potential), prec2double(q),
              prec2double (pot_mid), prec2double(pot_well - pot_mid), prec2double(eacc));
        }
      fflush (csvfp);

      if ( quick )
          continue;

      Screen screen(800, 800, 1.4, 1.4);
      if ( cross )
          screen.set_axis (vect3d(0, 0, 0), vect3d(1, 1, 0).normal(), vect3d(0, 0, 1));
      Statics statics;
      ElectroDynamics dyn(statics);

#if 1
      Octree tree;
      {
      printf ("pmap make\n");
      PMap pmap(vect3d(cage_radius + 0.01, cage_radius + 0.01, cage_radius + 0.01), 0.0075, 0, 0.015);
      printf ("pmap fill\n");
      pf.fill (pmap);
      printf ("pmap recalc\n");
      pmap.recalc();
      printf ("pmap done\n");
      //statics.addStaticElement (new StaticPMap (pmap));
      PolywellOctreeBuilder builder(pmap, vect3d(0, 0, 0), vect3d(0.8, 0.8, 0.8));
      builder.build (tree);
      }
      //statics.addStaticElement (new StaticOctree48 (tree, true, true, true, r, wr, spacing, 1, 1));
      statics.addStaticElement (new StaticOctree48 (tree, true, true, true, 1, 1));
#else
      statics.addStaticElement (new StaticPotentialField (pf));
#endif
      Scene scene(dyn);

      scene.calc(screen, Scene::CALC_POTENTIAL);
      scene.set_coloring (20, false);
      if ( i == LOOPS )
        {
          cali_pot_min = scene.get_fmin();
          cali_pot_max = scene.get_fmax();
          cali_pot_add = scene.get_fadd(); // potential can be negative
        }
      else
        {
          scene.calibrate_field (cali_pot_min, cali_pot_max, cali_pot_add);
        }
      scene.render_map (screen);
      screen.write (prefix + "_potential.bmp");

      scene.calc(screen, Scene::CALC_EFIELD);
      scene.set_coloring (0, false);
      if ( i == LOOPS )
        {
          cali_e_min = scene.get_fmin();
          cali_e_max = scene.get_fmax();
        }
      else
        {
          scene.calibrate_field (cali_e_min, cali_e_max);
        }
      scene.render_map (screen);
      screen.write (prefix + "_efield.bmp");
      scene.set_coloring (20, true);
      scene.render_map (screen);
      screen.write (prefix + "_efield_log_stepped.bmp");
      scene.set_coloring (0, true);
      scene.render_lic (screen);
      screen.write (prefix + "_efield_lic_logcolored.bmp");
      if ( snapshots )
        {
          Screen screen(380, 380, 1.4, 1.4);
          if ( cross )
              screen.set_axis (vect3d(0, 0, 0), vect3d(1, 1, 0).normal(), vect3d(0, 0, 1));
          Statics statics;
          ElectroDynamics dyn(statics);
          statics.addStaticElement (new StaticPotentialField (pf));
          Scene scene(dyn);

          scene.calc(screen, Scene::CALC_POTENTIAL);
          scene.set_coloring (20, false);
          scene.calibrate_field (cali_pot_min, cali_pot_max, cali_pot_add);
          scene.render_map (screen);
          screen.write (prefix + "_min_potential.bmp");

          scene.calc(screen, Scene::CALC_EFIELD);
          scene.set_coloring (0, false);
          scene.calibrate_field (cali_e_min, cali_e_max);
          scene.render_map (screen);
          screen.write (prefix + "_min_efield.bmp");
          scene.set_coloring (20, true);
          scene.render_map (screen);
          screen.write (prefix + "_min_efield_log_stepped.bmp");
          scene.set_coloring (0, true);
          scene.render_lic (screen);
          screen.write (prefix + "_min_efield_lic_logcolored.bmp");
        }
      break;
    }

  fclose (csvfp);

  return 0;
}

